#!/usr/bin/env python3
"""
PyTorch 环境配置模块

这个模块负责配置PyTorch环境，抑制常见的警告信息，
特别是NNPACK相关的警告，提供更清洁的运行环境。

使用方法:
    from kidsbuddy.utils.torch_config import configure_torch_environment
    configure_torch_environment()

或者在导入torch之前调用:
    from kidsbuddy.utils.torch_config import setup_torch_env_vars
    setup_torch_env_vars()
    import torch
"""

import os
import warnings
import logging
from typing import Optional


def setup_torch_env_vars() -> None:
    """
    设置PyTorch相关的环境变量
    
    这个函数应该在导入torch之前调用，以确保环境变量生效。
    """
    # 禁用NNPACK相关功能
    os.environ['NNPACK_DISABLE'] = '1'
    os.environ['PYTORCH_DISABLE_NNPACK_CHECK'] = '1'
    
    # 设置线程数以避免OpenMP和MKL警告
    os.environ['OMP_NUM_THREADS'] = '1'
    os.environ['MKL_NUM_THREADS'] = '1'
    os.environ['NUMEXPR_NUM_THREADS'] = '1'
    
    # 禁用一些可能导致警告的功能
    os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] = '1'  # macOS Metal Performance Shaders
    os.environ['PYTORCH_MPS_HIGH_WATERMARK_RATIO'] = '0.0'  # 禁用MPS内存管理警告


def configure_torch_warnings() -> None:
    """
    配置PyTorch相关的警告过滤器
    """
    # 抑制NNPACK相关警告
    warnings.filterwarnings("ignore", message=".*NNPACK.*")
    warnings.filterwarnings("ignore", message=".*Could not initialize NNPACK.*")
    warnings.filterwarnings("ignore", message=".*Unsupported hardware.*")
    
    # 抑制PyTorch用户警告
    warnings.filterwarnings("ignore", category=UserWarning, module="torch.*")
    warnings.filterwarnings("ignore", category=UserWarning, module="torchvision.*")
    
    # 抑制Ultralytics YOLO相关警告
    warnings.filterwarnings("ignore", category=UserWarning, module="ultralytics.*")
    warnings.filterwarnings("ignore", message=".*Ultralytics.*")
    
    # 抑制其他常见的机器学习库警告
    warnings.filterwarnings("ignore", category=FutureWarning, module="torch.*")
    warnings.filterwarnings("ignore", category=DeprecationWarning, module="torch.*")


def configure_torch_logging() -> None:
    """
    配置PyTorch相关的日志级别
    """
    # 设置PyTorch相关库的日志级别
    logging.getLogger('torch').setLevel(logging.ERROR)
    logging.getLogger('torchvision').setLevel(logging.ERROR)
    logging.getLogger('ultralytics').setLevel(logging.ERROR)
    
    # 抑制一些特定的日志记录器
    logging.getLogger('PIL').setLevel(logging.ERROR)
    logging.getLogger('matplotlib').setLevel(logging.ERROR)


def configure_torch_backend(device: Optional[str] = None) -> None:
    """
    配置PyTorch后端设置
    
    参数:
        device: 指定设备类型 ('cpu', 'cuda', 'mps' 等)，None表示自动选择
    """
    try:
        import torch
        
        # 设置线程数
        torch.set_num_threads(1)
        
        # 设置CPU优化选项
        try:
            torch.set_flush_denormal(True)
        except (AttributeError, RuntimeError):
            pass
        
        # 如果指定了设备，设置默认设备
        if device:
            try:
                if device == 'cuda' and torch.cuda.is_available():
                    torch.cuda.set_device(0)
                elif device == 'mps' and hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
                    # macOS Metal Performance Shaders
                    pass
            except (AttributeError, RuntimeError):
                pass
                
    except ImportError:
        # 如果torch还没有安装，跳过配置
        pass


def configure_torch_environment(
    device: Optional[str] = None,
    suppress_warnings: bool = True,
    configure_logging: bool = True
) -> None:
    """
    完整配置PyTorch环境
    
    这是主要的配置函数，会调用所有其他配置函数。
    
    参数:
        device: 指定设备类型，None表示自动选择
        suppress_warnings: 是否抑制警告
        configure_logging: 是否配置日志级别
    """
    # 设置环境变量（必须在导入torch之前）
    setup_torch_env_vars()
    
    # 配置警告
    if suppress_warnings:
        configure_torch_warnings()
    
    # 配置日志
    if configure_logging:
        configure_torch_logging()
    
    # 配置后端
    configure_torch_backend(device)


def get_optimal_device() -> str:
    """
    获取最优的计算设备
    
    返回:
        设备字符串 ('cuda', 'mps', 'cpu')
    """
    try:
        import torch
        
        # 优先使用CUDA
        if torch.cuda.is_available():
            return 'cuda'
        
        # 其次使用MPS (macOS)
        if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
            return 'mps'
        
        # 最后使用CPU
        return 'cpu'
        
    except ImportError:
        return 'cpu'


def print_torch_info() -> None:
    """
    打印PyTorch环境信息（用于调试）
    """
    try:
        import torch
        
        print("=== PyTorch 环境信息 ===")
        print(f"PyTorch 版本: {torch.__version__}")
        print(f"CUDA 可用: {torch.cuda.is_available()}")
        if torch.cuda.is_available():
            print(f"CUDA 版本: {torch.version.cuda}")
            print(f"GPU 数量: {torch.cuda.device_count()}")
        
        if hasattr(torch.backends, 'mps'):
            print(f"MPS 可用: {torch.backends.mps.is_available()}")
        
        print(f"推荐设备: {get_optimal_device()}")
        print(f"CPU 线程数: {torch.get_num_threads()}")
        
        # 检查环境变量
        env_vars = [
            'NNPACK_DISABLE',
            'PYTORCH_DISABLE_NNPACK_CHECK',
            'OMP_NUM_THREADS',
            'MKL_NUM_THREADS'
        ]
        
        print("\n=== 环境变量 ===")
        for var in env_vars:
            value = os.environ.get(var, '未设置')
            print(f"{var}: {value}")
            
    except ImportError:
        print("PyTorch 未安装")


# 自动配置（当模块被导入时）
if __name__ != "__main__":
    # 只在作为模块导入时自动配置环境变量
    setup_torch_env_vars()


if __name__ == "__main__":
    # 如果直接运行此脚本，显示环境信息
    configure_torch_environment()
    print_torch_info()
